from pathlib import Path


def classify_files_and_save_to_txt(directory, output_file=None):
    # 处理输出文件路径为空的情况
    if not output_file:
        # 获取当前脚本所在目录
        current_dir = Path(__file__).parent
        # 获取目录名称（不是完整路径）
        dir_name = Path(directory).name
        # 生成文件名：filetype+目录名称.txt
        output_filename = f"filetype_{dir_name}.txt"
        output_file = current_dir / output_filename
    else:
        output_file = Path(output_file)

    # 创建一个字典用于存储文件分类
    file_dict = {}

    # 使用 Path 对象遍历目录及其子目录
    for file in Path(directory).rglob('*'):
        if file.is_file():  # 确保是文件
            # 获取文件的后缀名
            file_suffix = file.suffix.lower()
            # 将文件名添加到相应分类的列表中
            if file_suffix not in file_dict:
                file_dict[file_suffix] = []
            file_dict[file_suffix].append(file.name)

    # 按文件后缀名排序
    sorted_suffixes = sorted(file_dict.keys())

    # 写入TXT文件
    with open(output_file, 'w', encoding='utf-8') as f:
        # 在首行添加目录路径描述
        f.write(f"以下是目录 '{directory}' 中的文件分类统计：\n\n")

        for suffix in sorted_suffixes:
            # 显示文件类型和数量
            f.write(f"文件类型: {suffix if suffix else '无扩展名'}，文件数量: {len(file_dict[suffix])}\n")
            # 列出文件
            for filename in sorted(file_dict[suffix]):
                f.write(f"  - {filename}\n")
            f.write("\n")  # 只保留空行分隔不同类型

    return output_file


if __name__ == "__main__":
    # 设置目录路径和输出文件名
    directory_path = r'G:\big_model\Voxtral-Mini-3B-2507'  # 替换为您的目录路径
    output_txt_file = None  # 设为None测试自动命名功能，也可指定具体路径

    # 调用函数进行文件分类
    saved_path = classify_files_and_save_to_txt(directory_path, output_txt_file)

    print(f"所有文件分类已保存到 {saved_path}")
